package com.itcast.peft.service;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

import org.apache.commons.io.IOUtils;
import org.apache.commons.text.StringSubstitutor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import com.itcast.peft.config.MultipleModelConfig.ChatModelHolder;
import com.itcast.peft.entity.Answer;
import com.itcast.peft.mapper.AnswerMapper;
import com.itcast.peft.mapper.QuestionMapper;
import com.itcast.peft.util.AsyncWrapper;
import com.itcast.peft.util.ExceptionWrapper;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

@Service
@Slf4j
public class EvaluateService {

    /**
     * 数据量较大，起并发
     */
    @Autowired
    private AsyncWrapper asyncWrapper;

    @Autowired
    private QuestionMapper questionMapper;

    @Autowired
    private AnswerMapper answerMapper;

    private ChatClient evalClient;
    private String evalTemplate;
    
    private ChatClient similarityEvalClient;
    private String similarityEvalTemplate;

    @Resource
    private ChatModelHolder distillEvalModelHolder;

    @Resource
    private ChatModelHolder userprefEvalModelHolder;
    
    @PostConstruct
    void init() throws IOException {
        try (var is = new ClassPathResource("prompt_eval.txt").getInputStream()) {
            this.evalTemplate = IOUtils.toString(is, StandardCharsets.UTF_8);
        }
        try (var is = new ClassPathResource("prompt_similarity_eval.txt").getInputStream()) {
            this.similarityEvalTemplate = IOUtils.toString(is, StandardCharsets.UTF_8);
        }
        this.evalClient = ChatClient.builder(distillEvalModelHolder.getModel()).build();
        this.similarityEvalClient = ChatClient.builder(userprefEvalModelHolder.getModel()).build();
    }

    @Async
    public void evaluate(String refGroup, String ansGroup) {
        var refMap = answerMapper.questionId2dataByGroup(refGroup);
        var ansMap = answerMapper.questionId2dataByGroup(ansGroup);

        var evalList = ansMap.keySet().stream().map(qid -> CompletableFuture
                .runAsync(() -> eval(refMap.get(qid), ansMap.get(qid), evalClient, evalTemplate), asyncWrapper.getExecutor())).toList();
        evalList.forEach(ExceptionWrapper.<CompletableFuture<Void>, Void>wrapToRTE(CompletableFuture::get)::apply);
        log.info("evaluation[{}] task completed:{} - {}", distillEvalModelHolder.getConfig().getModelName(), refGroup,
                ansGroup);
    }
    
    @Async
    public void similarityEvaluate(String refGroup, String ansGroup) {
        var refMap = answerMapper.questionId2dataByGroup(refGroup);
        var ansMap = answerMapper.questionId2dataByGroup(ansGroup);

        var evalList = ansMap.keySet().stream().limit(100).map(qid -> CompletableFuture
                .runAsync(() -> eval(refMap.get(qid), ansMap.get(qid), similarityEvalClient, similarityEvalTemplate), asyncWrapper.getExecutor())).toList();
        evalList.forEach(ExceptionWrapper.<CompletableFuture<Void>, Void>wrapToRTE(CompletableFuture::get)::apply);
        log.info("similarity evaluation[{}] task completed:{} - {}", userprefEvalModelHolder.getConfig().getModelName(), refGroup,
                ansGroup);
    }

    private void eval(Answer ref, Answer ans, ChatClient client,String template) {
        var question = questionMapper.selectById(ref.getQuestionId()).getQuestion();
        Map<String, String> params = new HashMap<>();

        params.put("question", question);
        params.put("answer", ans.getAnswer());
        params.put("ref_answer", ref.getAnswer());

        var userMessage = StringSubstitutor.replace(template, params, "${", "}");
        
        var opt = OpenAiChatOptions.builder()
                .temperature(0d)
                .build();
        var prompt = new Prompt(userMessage, opt);
        
        var eval = client.prompt(prompt).call().entity(EvalResult.class);

        ans.setEvalScore(eval.getScore());
        ans.setEvalResult(eval.getReason());

        var updateRes = answerMapper.updateById(ans);
        log.info("update single eval:{} on question:{}", updateRes, question);
    }
    
    @Data
    public static class EvalResult {
        String reason;
        Integer score;
    }

}
